package com.idp.saml;


import com.idp.metadata.IDPMetadata;
import com.idp.metadata.MetadataBean;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.opensaml.Configuration;
import org.opensaml.common.impl.RandomIdentifierGenerator;
import org.opensaml.saml2.core.*;
import org.opensaml.xml.XMLObjectBuilderFactory;
import org.opensaml.xml.schema.XSString;
import org.opensaml.xml.schema.impl.XSStringBuilder;

import javax.xml.namespace.QName;
import java.util.HashMap;
import java.util.Map;


public class SAMLResponse {
    //SAMLResponse 过期时间
    private static final int DEFAULT_ASSERTION_VALID_TIME_MIN = 5;
    private static final MetadataBean idpMetadata = IDPMetadata.getInstance();

    public static MetadataBean getIdpMetadata() {
        return idpMetadata;
    }

    public static Assertion buildAssertion(AuthnRequest authnRequest, MetadataBean spMetadata, Map<String, String> responseAssertMap) {
        String entityId = idpMetadata.getEntityId();
        Assertion assertion = buildSAMLObject(Assertion.class);
        DateTime notBefore = DateTime.now(DateTimeZone.UTC);
        DateTime notOnOrAfter = notBefore.plusMinutes(DEFAULT_ASSERTION_VALID_TIME_MIN);

        Subject subject = buildSubject(entityId, authnRequest.getAssertionConsumerServiceURL(), authnRequest.getID(), notOnOrAfter);

        Issuer issuer = buildIssuer(entityId);

        Audience audience = buildSAMLObject(Audience.class);
        audience.setAudienceURI(spMetadata.getEntityId());
        AudienceRestriction audienceRestriction = buildSAMLObject(AudienceRestriction.class);
        audienceRestriction.getAudiences().add(audience);

        Conditions conditions = buildSAMLObject(Conditions.class);
        conditions.getAudienceRestrictions().add(audienceRestriction);
        conditions.setNotBefore(notBefore);
        conditions.setNotOnOrAfter(notOnOrAfter);

        AuthnStatement authnStatement = buildAuthnStatement(new DateTime(), entityId);

        assertion.setID(new RandomIdentifierGenerator().generateIdentifier());
        assertion.setIssueInstant(new DateTime());
        assertion.setSubject(subject);
        assertion.setIssuer(issuer);
        assertion.setConditions(conditions);
        assertion.getAuthnStatements().add(authnStatement);
        Map<String, String> attribute = new HashMap<>();
        // 添加属性
        buildAttribute(attribute, responseAssertMap);
        assertion.getAttributeStatements().add(buildAttributeStatement(attribute));
        return assertion;
    }

    public static Response buildResponse(AuthnRequest authnRequest, Assertion assertion) {
        Response response = buildSAMLObject(Response.class);
        Issuer issuer = buildIssuer(idpMetadata.getEntityId());
        Status status = buildStatus(StatusCode.SUCCESS_URI);
        response.setID(new RandomIdentifierGenerator().generateIdentifier());
        response.setIssuer(issuer);
        response.setStatus(status);
        response.setIssueInstant(new DateTime());

        response.setInResponseTo(authnRequest.getID());
        response.setConsent(Response.UNSPECIFIED_CONSENT);
        response.setDestination(authnRequest.getAssertionConsumerServiceURL());
        response.getAssertions().add(assertion);
        return response;
    }

    private static Subject buildSubject(String subjectNameId, String recipient,
                                        String inResponseTo, DateTime notOnOrAfter) {
        Subject subject = buildSAMLObject(Subject.class);
        NameID nameID = buildSAMLObject(NameID.class);
        nameID.setValue(subjectNameId);
        nameID.setFormat(NameID.TRANSIENT);
        SubjectConfirmationData subjectConfirmationData = buildSAMLObject(SubjectConfirmationData.class);
        subjectConfirmationData.setRecipient(recipient);
        subjectConfirmationData.setNotOnOrAfter(notOnOrAfter);
        subjectConfirmationData.setInResponseTo(inResponseTo);
        SubjectConfirmation subjectConfirmation = buildSAMLObject(SubjectConfirmation.class);
        subjectConfirmation.setMethod(SubjectConfirmation.METHOD_BEARER);
        subjectConfirmation.setSubjectConfirmationData(subjectConfirmationData);
        subject.setNameID(nameID);
        subject.getSubjectConfirmations().add(subjectConfirmation);
        return subject;
    }

    private static AttributeStatement buildAttributeStatement(Map<String, String> attributes) {
        AttributeStatement attributeStatement = buildSAMLObject(AttributeStatement.class);
        for (Map.Entry<String, String> entry : attributes.entrySet()) {
            attributeStatement.getAttributes().add(buildAttribute(entry.getKey(), entry.getValue()));
        }
        return attributeStatement;
    }

    private static Issuer buildIssuer(String entityId) {
        Issuer issuer = buildSAMLObject(Issuer.class);
        issuer.setValue(entityId);
        issuer.setFormat(NameIDType.ENTITY);
        return issuer;
    }

    private static AuthnStatement buildAuthnStatement(DateTime authnInstant, String entityId) {
        AuthnContextClassRef authnContextClassRef = buildSAMLObject(AuthnContextClassRef.class);
        authnContextClassRef.setAuthnContextClassRef(AuthnContext.PASSWORD_AUTHN_CTX);

        AuthenticatingAuthority authenticatingAuthority = buildSAMLObject(AuthenticatingAuthority.class);
        authenticatingAuthority.setURI(entityId);

        AuthnContext authnContext = buildSAMLObject(AuthnContext.class);
        authnContext.setAuthnContextClassRef(authnContextClassRef);
        authnContext.getAuthenticatingAuthorities().add(authenticatingAuthority);

        AuthnStatement authnStatement = buildSAMLObject(AuthnStatement.class);
        authnStatement.setAuthnContext(authnContext);

        authnStatement.setAuthnInstant(authnInstant);
        return authnStatement;
    }

    private static Attribute buildAttribute(String name, String value) {
        XSStringBuilder stringBuilder = (XSStringBuilder) Configuration.getBuilderFactory().getBuilder(XSString.TYPE_NAME);

        Attribute attribute = buildSAMLObject(Attribute.class);
        attribute.setName(name);
        attribute.setNameFormat("urn:oasis:names:tc:SAML:2.0:attrname-format:uri");
        XSString stringValue = stringBuilder.buildObject(AttributeValue.DEFAULT_ELEMENT_NAME, XSString.TYPE_NAME);
        stringValue.setValue(value);
        attribute.getAttributeValues().add(stringValue);
        return attribute;
    }

    private static Status buildStatus(String value) {
        Status status = buildSAMLObject(Status.class);
        StatusCode statusCode = buildSAMLObject(StatusCode.class);

        statusCode.setValue(value);
        status.setStatusCode(statusCode);

        StatusMessage statusMessage = buildSAMLObject(StatusMessage.class);
        statusMessage.setMessage(value);
        status.setStatusMessage(statusMessage);
        return status;
    }

    //此处添加属性字段
    private static Map<String, String> buildAttribute(Map<String, String> attribute, Map<String, String> responseAssertMap) {
        attribute.put("xUserId", responseAssertMap.get("xUserId") == null ? "odder" : responseAssertMap.get("xUserId"));
        if (!responseAssertMap.isEmpty()) {
            for (Map.Entry<String, String> entry : responseAssertMap.entrySet()) {
                attribute.put(entry.getKey(), entry.getValue());
            }
        }
        return attribute;
    }

    public static <T> T buildSAMLObject(final Class<T> clazz) {
        T object = null;
        try {
            XMLObjectBuilderFactory builderFactory = Configuration.getBuilderFactory();
            QName defaultElementName = (QName) clazz.getDeclaredField("DEFAULT_ELEMENT_NAME").get(null);
            object = (T) builderFactory.getBuilder(defaultElementName).buildObject(defaultElementName);
        } catch (IllegalAccessException | NoSuchFieldException e) {
            throw new IllegalArgumentException("Could not create SAML object");
        }
        return object;
    }
}
